
import os
from tqdm import tqdm
import torch
import torch.nn.functional as F
from utils.utils import get_lr


#----------------------------------------------------------------------#
#   训练过程函数
#   输入:模型、损失函数、优化器、学习率下降策略、训练集、验证集、训练世代数、
#   多少世代保存一次权重、每个世代训练集的批次、每个世代验证集的批次、损失准确率可视化类
#----------------------------------------------------------------------#
def train(model, loss_func, optimizer, scheduler, train_loader, valid_loader, epoch_step, epoch_step_val, Epoch,
          save_period, save_dir,  loss_acc_history):

    # 是否使用GPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 将模型和损失函数放到device上, 即使用cpu还是使用gpu训练
    model.to(device)
    loss_func.to(device)

    #---------------#
    #   epochs循环
    #---------------#
    for epoch in range(Epoch):
        total_loss = 0
        total_accuracy = 0
        val_loss = 0
        val_accuracy = 0

        print('Start Train')
        pbar = tqdm(total=epoch_step, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
        #--------------------#
        #   训练数据batch循环
        #--------------------#
        #   训练过程中开启Dropout
        model.train()
        for iteration, batch in enumerate(train_loader):
            if iteration >= epoch_step:
                break
            #--------------------#
            # 获取图片和标签
            #--------------------#
            images, labels = batch
            #-------------------------------------------#
            # 将数据放到device上, 即使用cpu还是使用gpu训练
            #-------------------------------------------#
            images = images.to(device)
            labels = labels.to(device)
            #--------------------#
            # 将梯度清零
            #--------------------#
            optimizer.zero_grad()
            #--------------------#
            # 前向传播,计算模型输出
            #--------------------#
            model_outputs = model(images)
            #--------------------#
            # 计算损失函数
            #--------------------#
            loss = loss_func(model_outputs, labels)
            # 梯度反向传播
            loss.backward()
            # 执行一步权值更新
            optimizer.step()

            total_loss += loss.item()
            # 计算准确率
            with torch.no_grad():
                accuracy = torch.mean((torch.argmax(F.softmax(model_outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor))
                total_accuracy += accuracy.item()

            pbar.set_postfix(**{'total_loss': total_loss / (iteration + 1),
                                'accuracy'  : total_accuracy / (iteration + 1),
                                'lr'        : get_lr(optimizer)})
            pbar.update(1)

        pbar.close()
        print('Finish Train')

        print('Start Validation')
        pbar = tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
        #--------------------#
        #   验证数据batch循环
        #--------------------#
        #   验证过程中关闭Dropout
        model.eval()
        for iteration, batch in enumerate(valid_loader):
            if iteration >= epoch_step_val:
                break
            # 获取图片和标签
            images, labels = batch
            with torch.no_grad():
                # 将数据放到device上, 即使用cpu还是使用gpu训练
                images = images.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()
                # 计算模型输出
                model_outputs = model(images)
                # 计算损失函数
                loss = loss_func(model_outputs, labels)
                val_loss += loss.item()
                # 计算准确率
                accuracy = torch.mean((torch.argmax(F.softmax(model_outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor))
                val_accuracy += accuracy.item()

            pbar.set_postfix(**{'total_loss': val_loss / (iteration + 1),
                                'accuracy'  : val_accuracy / (iteration + 1),
                                'lr'        : get_lr(optimizer)})
            pbar.update(1)

        pbar.close()
        print('Finish Validation')
        loss_acc_history.append_loss(epoch, total_loss / epoch_step, val_loss / epoch_step_val)
        loss_acc_history.append_acc(epoch, total_accuracy / epoch_step, val_accuracy / epoch_step_val)
        loss_acc_history.append_lr(epoch, get_lr(optimizer))
        print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
        print('Total Loss: %.3f || Val Loss: %.3f ' % (total_loss / epoch_step, val_loss / epoch_step_val))

        #-----------------------------------------------#
        #   保存权值
        #-----------------------------------------------#
        if (epoch + 1) % save_period == 0 or epoch + 1 == Epoch:
            # 保存模型参数, 仅保存模型参数
            torch.save(model.state_dict(), os.path.join(save_dir, "ep%03d-loss%.3f-val_loss%.3f.pth" % (epoch + 1, total_loss / epoch_step, val_loss / epoch_step_val)))

        if len(loss_acc_history.val_loss) <= 1 or (val_loss / epoch_step_val) <= min(loss_acc_history.val_loss):
            print('Save best model to best_epoch_weights.pth')
            torch.save(model.state_dict(), os.path.join(save_dir, "best_epoch_weights.pth"))

        torch.save(model.state_dict(), os.path.join(save_dir, "last_epoch_weights.pth"))

        # 更新学习率
        scheduler.step()

#----------------------------------------------------------------------#
#   训练过程函数
#   输入:模型、损失函数、优化器、学习率下降策略、训练集、验证集、训练世代数、
#   多少世代保存一次权重、每个世代训练集的批次、每个世代验证集的批次、损失准确率可视化类
#----------------------------------------------------------------------#
def lm_softmax_train(model, loss_func, optimizer, scheduler, train_loader, valid_loader, epoch_step, epoch_step_val, Epoch,
          save_period, save_dir,  loss_acc_history):

    # 是否使用GPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 将模型和损失函数放到device上, 即使用cpu还是使用gpu训练
    model.to(device)
    loss_func.to(device)

    #---------------#
    #   epochs循环
    #---------------#
    for epoch in range(Epoch):
        total_loss = 0
        total_accuracy = 0
        val_loss = 0
        val_accuracy = 0

        print('Start Train')
        pbar = tqdm(total=epoch_step, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
        #--------------------#
        #   训练数据batch循环
        #--------------------#
        #   训练过程中开启Dropout
        model.train()
        for iteration, batch in enumerate(train_loader):
            if iteration >= epoch_step:
                break
            #--------------------#
            # 获取图片和标签
            #--------------------#
            images, labels = batch
            #-------------------------------------------#
            # 将数据放到device上, 即使用cpu还是使用gpu训练
            #-------------------------------------------#
            images = images.to(device)
            labels = labels.to(device)
            #--------------------#
            # 将梯度清零
            #--------------------#
            optimizer.zero_grad()
            #--------------------#
            # 前向传播,计算模型输出
            #--------------------#
            model_outputs, _ = model(images, labels)
            #--------------------#
            # 计算损失函数
            #--------------------#
            loss = loss_func(model_outputs, labels)
            # 梯度反向传播
            loss.backward()
            # 执行一步权值更新
            optimizer.step()

            total_loss += loss.item()
            # 计算准确率
            with torch.no_grad():
                accuracy = torch.mean((torch.argmax(F.softmax(model_outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor))
                total_accuracy += accuracy.item()

            pbar.set_postfix(**{'total_loss': total_loss / (iteration + 1),
                                'accuracy'  : total_accuracy / (iteration + 1),
                                'lr'        : get_lr(optimizer)})
            pbar.update(1)

        pbar.close()
        print('Finish Train')

        print('Start Validation')
        pbar = tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)
        #--------------------#
        #   验证数据batch循环
        #--------------------#
        #   验证过程中关闭Dropout
        model.eval()
        for iteration, batch in enumerate(valid_loader):
            if iteration >= epoch_step_val:
                break
            # 获取图片和标签
            images, labels = batch
            with torch.no_grad():
                # 将数据放到device上, 即使用cpu还是使用gpu训练
                images = images.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()
                # 计算模型输出
                model_outputs, _ = model(images)
                # 计算损失函数
                loss = loss_func(model_outputs, labels)
                val_loss += loss.item()
                # 计算准确率
                accuracy = torch.mean((torch.argmax(F.softmax(model_outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor))
                val_accuracy += accuracy.item()

            pbar.set_postfix(**{'total_loss': val_loss / (iteration + 1),
                                'accuracy'  : val_accuracy / (iteration + 1),
                                'lr'        : get_lr(optimizer)})
            pbar.update(1)

        pbar.close()
        print('Finish Validation')
        loss_acc_history.append_loss(epoch, total_loss / epoch_step, val_loss / epoch_step_val)
        loss_acc_history.append_acc(epoch, total_accuracy / epoch_step, val_accuracy / epoch_step_val)
        loss_acc_history.append_lr(epoch, get_lr(optimizer))
        print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
        print('Total Loss: %.3f || Val Loss: %.3f ' % (total_loss / epoch_step, val_loss / epoch_step_val))

        #-----------------------------------------------#
        #   保存权值
        #-----------------------------------------------#
        if (epoch + 1) % save_period == 0 or epoch + 1 == Epoch:
            # 保存模型参数, 仅保存模型参数
            torch.save(model.state_dict(), os.path.join(save_dir, "ep%03d-loss%.3f-val_loss%.3f.pth" % (epoch + 1, total_loss / epoch_step, val_loss / epoch_step_val)))

        if len(loss_acc_history.val_loss) <= 1 or (val_loss / epoch_step_val) <= min(loss_acc_history.val_loss):
            print('Save best model to best_epoch_weights.pth')
            torch.save(model.state_dict(), os.path.join(save_dir, "best_epoch_weights.pth"))

        torch.save(model.state_dict(), os.path.join(save_dir, "last_epoch_weights.pth"))

        # 更新学习率
        scheduler.step()


